package com.wql;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Properties;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.transform.OutputKeys;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;

import org.mybatis.generator.api.MyBatisGenerator;
import org.mybatis.generator.config.Configuration;
import org.mybatis.generator.config.xml.ConfigurationParser;
import org.mybatis.generator.internal.DefaultShellCallback;
import org.w3c.dom.Document;
import org.w3c.dom.DocumentType;
import org.w3c.dom.NamedNodeMap;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

public class Main {
	private static String generatorConfig = "generatorConfig.xml";
	private String configPropertie = "config.properties";
	private boolean needConvert = false;

	private Properties properties = new Properties();
	private List<String> tables = new ArrayList<String>();

	public static void main(String[] args) throws Exception {
		InputStream inputStream = Main.class.getClassLoader().getResourceAsStream(generatorConfig);
		// 一次生成多张表时使用
		inputStream = new Main().generatorManyTablesIfNeed(inputStream);

		// 正常生成器逻辑
		boolean overwrite = true;
		List<String> warnings = new ArrayList<String>();
		ConfigurationParser cp = new ConfigurationParser(warnings);
		Configuration config = cp.parseConfiguration(inputStream);
		DefaultShellCallback callback = new DefaultShellCallback(overwrite);
		MyBatisGenerator myBatisGenerator = new MyBatisGenerator(config, callback, warnings);
		myBatisGenerator.generate(null);
		System.out.println("----ok----");
	}

	// 需要生成很多表时，需要在generatorConfig中配置多个table标签，很繁琐。
	// 这里通过在属性文件中tableName中配置通配符，多个表名，甚至正则表达式来匹配数据库中所有表。
	// 例如：tableName=*(匹配数据库中所有表)
	// 例如：tableName=t1,t2,t3或者tableName=t1 t2 t3(匹配数据中t1,t2,t3表)
	// 例如：tableName=prod.*(匹配数据库中prod打头的表)
	private InputStream generatorManyTablesIfNeed(InputStream inputStream) {
		return init().getFiltedTablesIfNeed().convertInputStream(inputStream);
	}

	// 加载属性配置文件，判断是否需要转换
	private Main init() {
		try {
			properties.load(Main.class.getClassLoader().getResourceAsStream(configPropertie));
			String tableNameConfig = properties.getProperty("tableName", "");
			needConvert = !tableNameConfig.matches("^[a-zA-Z0-9_-]+$");
		} catch (IOException e) {
			e.printStackTrace();
		}
		return this;
	}

	// 使用属性文件中配置的表名过滤规则过滤数据库中所有表
	private Main getFiltedTablesIfNeed() {
		if (needConvert) {
			try {
				String tableNameStr = properties.getProperty("tableName"),
						driverStr = properties.getProperty("jdbc.driver"), urlStr = properties.getProperty("jdbc.url"),
						usernameStr = properties.getProperty("jdbc.username"),
						passwordStr = properties.getProperty("jdbc.password");
				List<String> tableNames = new ArrayList<String>();
				Class.forName(driverStr);
				Connection connection = DriverManager.getConnection(urlStr, usernameStr, passwordStr);
				DatabaseMetaData dbMetData = connection.getMetaData();
				ResultSet rs = dbMetData.getTables(null, null, null, new String[] { "TABLE", "VIEW" });
				if (tableNameStr.contains(" ") || tableNameStr.contains(",")) {
					tableNames.addAll(Arrays.asList(tableNameStr.split("\\s+|,")));
					tableNameStr = "";
				}
				while (rs.next()) {
					if (rs.getString(4) != null && (rs.getString(4).equalsIgnoreCase("TABLE")
							|| rs.getString(4).equalsIgnoreCase("VIEW"))) {
						String name = rs.getString(3).toLowerCase();
						if (!tableNames.isEmpty()) {
							for (String tableName : tableNames) {
								if (name.matches(tableName)) {
									tables.add(name);
								}
							}
						} else {
							if ("*".equals(tableNameStr) || name.matches(tableNameStr)) {
								tables.add(name);
							}
						}
					}
				}
				System.out.println("filter is: " + tableNameStr + tableNames);
				System.out.println("filted tables: " + tables.toString());
			} catch (Exception e) {
				e.printStackTrace();
			}
		}
		return this;
	}

	// 通过运行时修改generatorConfig文件，增加table标签达到生成多张表的功能
	private InputStream convertInputStream(InputStream inputStream) {
		if (needConvert && tables.size() > 0) {
			try {
				DocumentBuilder builder = DocumentBuilderFactory.newInstance().newDocumentBuilder();
				Document document = builder.parse(inputStream);
				if (null != document) {
					NodeList nodes = document.getElementsByTagName("properties");
					if (nodes.getLength() == 1) {
						NamedNodeMap attributes = nodes.item(0).getAttributes();
						attributes.getNamedItem("resource").setNodeValue(configPropertie);
					}
					nodes = document.getElementsByTagName("table");
					if (nodes.getLength() == 1) {
						Node tableNode = nodes.item(0);
						Node parentNode = tableNode.getParentNode();
						for (String tableName : tables) {
							Node newTable = tableNode.cloneNode(false);
							NamedNodeMap attributes = newTable.getAttributes();
							attributes.getNamedItem("tableName").setNodeValue(tableName);
							parentNode.appendChild(newTable);
						}
						parentNode.removeChild(tableNode);
						ByteArrayOutputStream bos = new ByteArrayOutputStream();
						TransformerFactory transformerFactory = TransformerFactory.newInstance();
						DocumentType doctype = document.getDoctype();
						Transformer transformer = transformerFactory.newTransformer();
						transformer.setOutputProperty(OutputKeys.DOCTYPE_PUBLIC, doctype.getPublicId());
						transformer.setOutputProperty(OutputKeys.DOCTYPE_SYSTEM, doctype.getSystemId());
						transformer.setOutputProperty(OutputKeys.ENCODING, document.getXmlEncoding());
						transformer.setOutputProperty(OutputKeys.METHOD, "xml");
						transformer.transform(new DOMSource(document), new StreamResult(bos));
						return new ByteArrayInputStream(bos.toByteArray());
					}
				}
			} catch (Exception e) {
				e.printStackTrace();
			}
		}
		return inputStream;
	}

}
